import cv2
import numpy as np
import mediapipe as mp

def draw_skeleton(image, landmarks):
    """
    根据姿势关键点绘制骨架图。
    """
    # 创建空白图像，用于绘制骨架
    skeleton_image = np.ones_like(image) * 255

    # 头部（实心圆）
    head = (int(landmarks[0].x * image.shape[1]), int(landmarks[0].y * image.shape[0]))
    cv2.circle(skeleton_image, head, 30, (255, 255, 255,255), -1)

    # 躯干（实心椭圆）
    shoulder = (int(landmarks[11].x * image.shape[1]), int(landmarks[11].y * image.shape[0]))
    pelvis = (int(landmarks[23].x * image.shape[1]), int(landmarks[23].y * image.shape[0]))
    center = ((shoulder[0] + pelvis[0]) // 2, (shoulder[1] + pelvis[1]) // 2)
    axes = (30, 70)  # 椭圆长短轴
    cv2.ellipse(skeleton_image, center, axes, 0, 0, 360, (0, 0, 0), -1)

    # 四肢（线条）
    keypoints = [
        (11, 13),  # 左上臂
        (13, 15),  # 左前臂
        (12, 14),  # 右上臂
        (14, 16),  # 右前臂
        (23, 25),  # 左大腿
        (25, 27),  # 左小腿
        (24, 26),  # 右大腿
        (26, 28),  # 右小腿
    ]
    for p1, p2 in keypoints:
        pt1 = (int(landmarks[p1].x * image.shape[1]), int(landmarks[p1].y * image.shape[0]))
        pt2 = (int(landmarks[p2].x * image.shape[1]), int(landmarks[p2].y * image.shape[0]))
        cv2.line(skeleton_image, pt1, pt2, (0, 0, 0), 10)

    return skeleton_image

# 加载图片
input_image_path = './zishi.png'  # 替换为您的图片路径
output_image_path = './skeleton_output.jpg'

# 初始化 MediaPipe
mp_pose = mp.solutions.pose
pose = mp_pose.Pose(static_image_mode=True, min_detection_confidence=0.5)

# 读取图片
image = cv2.imread(input_image_path)
if image is None:
    raise FileNotFoundError(f"Image not found at {input_image_path}")

# 转换为 RGB
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 获取姿势关键点
results = pose.process(image_rgb)
if not results.pose_landmarks:
    raise ValueError("No pose landmarks detected!")

# 绘制骨架图
skeleton_image = draw_skeleton(image, results.pose_landmarks.landmark)

# 保存结果
cv2.imwrite(output_image_path, skeleton_image)
cv2.imshow("huizhi",skeleton_image)
cv2.waitKey(0)
cv2.destroyAllWindows()
print(f"Skeleton image saved to {output_image_path}")